source(here::here("R/00_setup.R"))

d_baseline = read_csv(here("data/precinct_baseline.csv"), show_col_types=FALSE)
model = read_rds(here("data/election_model.rds"))
hex = read_rds(here("data/districts_2020_hex.rds")) |>
  rename(cd_2020 = district)
hex_states = group_by(hex, state) |>
  summarize(is_coverage=TRUE)

natl_baseline = with(d_baseline, log(sum(ndv)) - log(sum(nrv)))

# matrix of logit D vote shares
m_pr_win = do.call(rbind, map(cli_progress_along(state.abb), function(i) {
  abbr = state.abb[i]
  d_st = filter(d_baseline, state == abbr)
  plans = load_50state_plans(abbr)
  m_out = part_dvs(plans, d_st, ndv, nrv) |>
    model$prob_win_natl() |>
    matrix(ncol=5001)
  rownames(m_out) = rep(abbr, nrow(m_out))
  m_out
}))


d_gerry = map_dfr(cli_progress_along(state.abb), function(i) {
  abbr = state.abb[i]
  plans = load_50state_plans(abbr)
  map = load_50state_map(abbr)
  plans$pr_dem = as.numeric(m_pr_win[rownames(m_pr_win) == abbr, ])
  pr_samp = avg_by_prec(plans, pr_dem, draws=NA)
  pr_enac = avg_by_prec(plans, pr_dem, draws="cd_2020")
  as_tibble(map) |>
    mutate(pr_enac = pr_enac,
           diff = pr_samp - pr_enac) |>
    group_by(state, cd_2020) |>
    summarize(pr_win = pr_enac[1],
              gerry = weighted.mean(diff, ndv+nrv))
})

# Figures -------


left_join(hex, d_gerry, by=c("state", "cd_2020")) |>
  group_by(state) |>
  mutate(cd_2020 = if_else(rep(n(), n()) > 1, as.character(cd_2020), state)) |>
  ggplot(aes(fill=gerry)) +
  geom_sf(lwd=0.15, color="black") +
  geom_sf(data=hex_states, inherit.aes=FALSE, lwd=0.4, color="black", fill=NA) +
  geom_sf_text(aes(geometry=geom_label, label=cd_2020),
               size=1.8, fontface="bold", color="#000000aa") +
  scale_fill_party_c("Change in win probability\nversus simulations", midpoint=0,
                     limits=c(-0.5, 0.5), reverse=TRUE,
                     labels=label_party_margin(midpoint=0.0, reverse=TRUE)) +
  ggplot2::theme_void(base_family="Arial", base_size=8) +
  coord_sf(expand=FALSE) +
  theme(plot.margin=margin(2, 2, 2, 2),
        legend.key.width=unit(0.5, "cm"),
        legend.direction="vertical",
        legend.position=c(0.9, 0.4))


ggsave(here("paper/figures/manipulation.pdf"), width=10, height=7.25, device=cairo_pdf)


left_join(hex, d_gerry, by=c("state", "cd_2020")) |>
  group_by(state) |>
  mutate(cd_2020 = if_else(rep(n(), n()) > 1, as.character(cd_2020), state)) |>
ggplot(aes(fill=pr_win)) +
  geom_sf(lwd=0.15, color="black") +
  geom_sf(data=hex_states, inherit.aes=FALSE, lwd=0.4, color="black", fill=NA) +
  geom_sf_text(aes(geometry=geom_label, label=cd_2020),
               size=1.8, fontface="bold", color="#000000aa") +
  scale_fill_party_c("Win probability") +
  ggplot2::theme_void(base_family="Arial", base_size=8) +
  coord_sf(expand=FALSE) +
  theme(plot.margin=margin(2, 2, 2, 2),
        legend.position=c(0.875, 0.375))

ggsave(here("paper/figures/pr_win.pdf"), width=10, height=7.25, device=cairo_pdf)



# Schematic for appendix ------------
plans <- load_50state_plans("NC")
map <- load_50state_map("NC")
plans$pr_dem <- as.numeric(m_pr_win[rownames(m_pr_win) == "NC", ])
d_nc <- as_tibble(map) |>
  transmute(
    cd_2020 = cd_2020,
    pr_samp = avg_by_prec(plans, pr_dem, draws = NA),
    pr_enac = avg_by_prec(plans, pr_dem, draws = "cd_2020"),
    geometry = rmapshaper::ms_simplify(geometry, 0.001, keep_shapes = TRUE)
  ) |>
  st_as_sf()

p2a <- ggplot(d_nc) +
  geom_sf(aes(fill = 1 - pr_enac), lwd = 0, color = NA) +
  geom_sf_text(
    data = d_nc |> group_by(cd_2020) |> summarize(),
    aes(label = cd_2020), fun.geometry = \(x) st_circle_center(st_as_sf(x)),
    size = 2, fontface = "bold", color = "#000000aa"
  ) +
  geom_district(aes(group = cd_2020), fill = NA, size = 0.3) +
  annotate("text", label = "(a) Enacted plan", x = 300e3, y = 50e3, size = 2.5) +
  scale_fill_wa_c(
    name = "Prob. represented\nby Republican", palette = "stuart",
    midpoint = 0.5, labels = scales::percent, limits = c(0, 1),
    reverse = TRUE
  ) +
  guides(fill = "none") +
  theme_void(base_family = "Arial", base_size = 8) +
  theme(
    plot.caption = element_text(size = 8, family = "Arial", hjust = 0.1, vjust = 10.0),
    plot.caption.position = "panel"
  )

p2b <- ggplot(d_nc) +
  geom_sf(aes(fill = 1 - pr_samp), lwd = 0, color = NA) +
  geom_sf_text(
    data = d_nc |> group_by(cd_2020) |> summarize(),
    aes(label = cd_2020), fun.geometry = \(x) st_circle_center(st_as_sf(x)),
    size = 2, fontface = "bold", color = "#000000aa"
  ) +
  geom_district(aes(group = cd_2020), fill = NA, size = 0.3) +
  scale_fill_wa_c(
    name = "Prob. represented\nby Republican", palette = "stuart",
    midpoint = 0.5, labels = scales::percent, limits = c(0, 1),
    reverse = TRUE
  ) +
  annotate("text", label = "(b) Average of\nsampled plans", x = 300e3, y = 70e3, size = 2.5) +
  theme_void(base_family = "Arial", base_size = 6) +
  theme(
    legend.position = "right",
    legend.direction = "vertical",
    plot.margin = margin()
  )

p2c <- ggplot(d_nc) +
  geom_sf(aes(fill = pr_samp - pr_enac), lwd = 0, color = NA) +
  geom_sf_text(
    data = d_nc |> group_by(cd_2020) |> summarize(),
    aes(label = cd_2020), fun.geometry = \(x) st_circle_center(st_as_sf(x)),
    size = 2, fontface = "bold", color = "#000000aa"
  ) +
  geom_district(aes(group = cd_2020), fill = NA, size = 0.3) +
  scale_fill_party_c("Change in win probability\nversus simulations\n",
    midpoint = 0,
    limits = c(-0.5, 0.5), reverse = TRUE,
    labels = label_party_margin(midpoint = 0.0, reverse = TRUE)
  ) +
  annotate("text", label = "(c) Change in win probability", x = 300e3, y = 70e3, size = 2.5) +
  theme_void(base_family = "Arial", base_size = 6) +
  guides(fill = guide_colorbar(title.position = "left", title.hjust = 1)) +
  theme(
    legend.position = "none",
    legend.direction = "vertical",
    plot.margin = margin()
  )


p2d <- left_join(hex, d_gerry, by = c("state", "cd_2020")) |>
  filter(state == "NC") |>
  rotate_shp(rot_rads = pi / 20) |>
  ggplot(aes(fill = gerry)) +
  geom_sf(lwd = 0.15, color = "black") +
  geom_sf(
    data = hex_states |> filter(state == "NC") |> rotate_shp(rot_rads = pi / 20),
    inherit.aes = FALSE, lwd = 0.4, color = "black", fill = NA
  ) +
  geom_sf_text(aes(label = cd_2020),
    size = 1.8, fontface = "bold", color = "#000000aa"
  ) +
  scale_fill_party_c("Change in win probability\nversus simulations",
    midpoint = 0,
    limits = c(-0.5, 0.5), reverse = TRUE,
    labels = label_party_margin(midpoint = 0.0, reverse = TRUE)
  ) +
  ggplot2::theme_void(base_family = "Arial", base_size = 6) +
  annotate("text", label = "(d) Aggregate change in win probability", x = 1200000, y = -600000, size = 2.5) +
  theme(
    plot.margin = margin(2, 2, 2, 2),
    legend.direction = "vertical",
    legend.position = "right"
  )

p2a + p2b + p2c + p2d
ggsave(here("paper/figures/nc_schem.pdf"), width = 10, height = 4.5, device = cairo_pdf)
